import tensorflow as tf
import numpy as np

X = tf.Variable([[[1, 0, 3], [4, 5, 6]],[[1, 2, 3], [4, 6, 6]]], dtype=tf.float32);

max_0 = tf.reduce_max(X, axis=0)
max_1 = tf.reduce_max(X, axis=1)
max_0_1 = tf.reduce_max(X, axis=[1, 0])
max_1_0 = tf.reduce_max(X, axis=[0, 1])
 
d_max_0 = tf.gradients(max_0, X)
d_max_1 = tf.gradients(max_1, X)
d_max_0_1 = tf.gradients(max_0_1, X)
d_max_1_0 = tf.gradients(max_1_0, X)
 
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
     
    print('------max_0------')
    print(sess.run(max_0))
    print('------max_1------')
    print(sess.run(max_1))
    print('------max_0_1------')
    print(sess.run(max_0_1))
    print('------max_1_0------')
    print(sess.run(max_1_0))
      
    print('------d_max_0------')
    print(sess.run(d_max_0))
    print('------d_max_1------')
    print(sess.run(d_max_1))
    print('------d_max_0_1------')
    print(sess.run(d_max_0_1))
    print('------d_max_1_0------')
    print(sess.run(d_max_1_0))